# This example demonstrates working with MongoDB Atlas

import os
from llmware.configs import LLMWareConfig
from llmware.library import Library
from llmware.retrieval import Query
from llmware.setup import Setup

'''
You can use Mongo Atlas in 2 ways:

1. Only as a vector DB (and use local or another mongo for data storage)
   In this case, you can set the following environment variable:
      MONGO_ATLAS_CONNECTION_URI=mongo+srv://<username>:<password>@<cluster-domain>...

2. For both vector and data storage
   In this case, you can set the following environment configuration:
      LLMWareConfig().set_config("collection_db_uri", "mongo+srv://<username>:<password>@<cluster-domain>..."

This example demonstrates the 2nd approach
'''

# Set this env var appropriately or just paste in your Mongo Atlas connection string:
mongo_atlas_connection_string = os.environ["MONGO_ATLAS_CONNECTION_URI"]

# Use MongoDB Atlas for both data storage and vector embedding storage/search
def use_mongo_atlas():

      # Ensure the global config for data storage is set to Mongo Atlas     
      LLMWareConfig().set_config("collection_db_uri",mongo_atlas_connection_string)

      # Create a library and populate it with some sample documents
      library_name="test_mongo_atlas"
      print (f"\n > Creating library {library_name}...")
      library = Library().create_new_library(library_name)
      print (f"\n > Loading the llmware Sample Files...")
      sample_files_path = Setup().load_sample_files()
      print (f"\n > Adding some files to the library...")
      library.add_files(input_folder_path=os.path.join(sample_files_path, "SmallLibrary"))

      # Create vector embedings using Mongo Atlas
      print (f"\n > Generating embedding vectors (using the 'mini-lm-sbert' model) and storing them (in 'Mongo Atlas')...")
      library.install_new_embedding(embedding_model_name="mini-lm-sbert", vector_db="mongo_atlas")

      # Do a semantic search in the library
      print (f"\n > Running a query for 'Salary'...")
      query_results = Query(library).semantic_query(query="salary", result_count=3, results_only=True)
      print (f"\n\nHits for 'Salary' in {library_name}:\n")
      for query_result in query_results:
            print("File: " +  query_result["file_source"] + " (Page " + str(query_result["page_num"]) + "):\n" + query_result["text"] + "\n")

use_mongo_atlas()